package net.cyue.ort.llm.data;

import ai.onnxruntime.OnnxTensor;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;

/**
 * Beam搜索树
 */
public class BeamSearchTree {
    
    public static class BeamSearchBranch {
        private final TokenList tokenList;
        private final Map<String, OnnxTensor> cache;
        
        public BeamSearchBranch(TokenList tokenList, Map<String, OnnxTensor> cache) {
            this.tokenList = tokenList;
            this.cache = cache;
        }
        
        public BeamSearchBranch(List<Token> tokenList, Map<String, OnnxTensor> cache) {
            this.tokenList = new TokenList(tokenList);
            this.cache = cache;
        }
        
        public long[] getTokenIDs() {
            return tokenList.getTokenIDs();
        }
        
        public List<Long> getTokenIDList() {
            return tokenList.getTokenIDList();
        }
        
        public BeamCandidate toBeamCandidate() {
            return new BeamCandidate(
                tokenList,
                tokenList.getSequenceScore(),
                this.cache
            );
        }
        
        public TokenList getTokenList() {
            return tokenList;
        }
    }
    
    public static class BeamSearchNode {
        private final TokenList tokenList;
        private final Map<String, OnnxTensor> cache;
        private final List<BeamSearchNode> children = new ArrayList<>();
        BeamSearchNode parent;
        
        public BeamSearchNode(
            BeamSearchNode parent,
            Token token,
            Map<String, OnnxTensor> cache
        ) {
            this.parent = parent;
            this.tokenList = new TokenList();
            this.cache = cache;
            if (token != null) {
                this.tokenList.add(token);
            }
        }
        
        public BeamSearchNode(
            BeamSearchNode parent,
            TokenList tokenList,
            Map<String, OnnxTensor> cache
        ) {
            this.parent = parent;
            this.tokenList = tokenList != null ? tokenList : new TokenList();
            this.cache = cache;
        }
        
        public BeamSearchNode(
            BeamSearchNode parent,
            List<Token> tokenList,
            Map<String, OnnxTensor> cache
        ) {
            this.parent = parent;
            this.tokenList = new TokenList(tokenList);
            this.cache = cache;
        }
        
        public BeamCandidate toBeamCandidate() {
            return new BeamCandidate(
                tokenList,
                tokenList.getSequenceScore(),
                cache
            );
        }
        
        public Map<String, OnnxTensor> getCache() {
            return cache;
        }
        
        public void addChild(BeamSearchNode child) {
            if (child.parent != null && !child.parent.equals(this)) {
                child.parent.children.remove(child);
            }
            child.parent = this;
            this.children.add(child);
        }
        
        public BeamSearchNode getChild(int index) {
            return (index >= 0 && index < children.size()) ? children.get(index) : null;
        }
        
        public long[] getTokenIDs() {
            if (tokenList == null || tokenList.isEmpty()) {
                return null;
            }
            return tokenList.getTokenIDs();
        }
        
        public List<Long> getTokenIDList() {
            if (tokenList == null || tokenList.isEmpty()) {
                return null;
            }
            return tokenList.getTokenIDList();
        }
        
        public List<Long> getParentIDList(boolean includeSelf) {
            List<Long> idList = new ArrayList<>();
            BeamSearchNode node;
            if (includeSelf) {
                node = this;
            } else {
                node = this.parent;
            }
            while (node != null) {
                List<Long> nodeIDList = node.getTokenIDList();
                if (nodeIDList == null || nodeIDList.isEmpty()) {
                    node = node.parent;
                    continue;
                }
                List<Long> reverseIDList = new ArrayList<>(nodeIDList.size());
                for (int i = nodeIDList.size() - 1; i >= 0; i--) {
                    reverseIDList.add(nodeIDList.get(i));
                }
                idList.addAll(reverseIDList);
                node = node.parent;
            }
            Collections.reverse(idList);
            return idList;
        }
        
        public List<BeamSearchBranch> getAllBranches() {
            return getAllBranches(true);
        }
        
        public List<BeamSearchBranch> getAllBranches(boolean includeSelf) {
            List<BeamSearchBranch> allBranches = new ArrayList<>();
            TokenList tokenList = this.tokenList;
            if (!includeSelf) {
                tokenList = new TokenList();
            }
            if (this.children.isEmpty()) {
                // 如果当前节点是叶节点,直接用当前token列表构造BeamSearchBranch
                BeamSearchBranch branch = new BeamSearchBranch(
                    tokenList,
                    this.cache
                );
                allBranches.add(branch);
            } else {
                // 对于每个子节点，获取其所有分支，将当前token列表插入到子分支token列表的头部
                for (BeamSearchNode child : this.children) {
                    for (BeamSearchBranch childBranch : child.getAllBranches()) {
                        List<Token> branchTokens = new ArrayList<>();
                        for (Token token : tokenList) {
                            branchTokens.add(token);
                        }
                        for (Token token : childBranch.getTokenList()) {
                            branchTokens.add(token);
                        }
                        BeamSearchBranch newBranch = new BeamSearchBranch(
                            branchTokens,
                            childBranch.cache
                        );
                        allBranches.add(newBranch);
                    }
                }
            }
            return allBranches;
        }
        
        public double getScore() {
            return tokenList.getSequenceScore();
        }
        
        public int getDepth() {
            int depth = 0;
            BeamSearchNode node = this;
            while (node.parent != null) {
                depth++;
                node = node.parent;
            }
            return depth;
        }
    }
    
    private final BeamSearchNode head;
    
    public BeamSearchTree(BeamSearchNode head) {
        this.head = head;
    }
    
    public BeamSearchNode getHead() {
        return head;
    }
    
    public int getDepth() {
        int depth = 1;
        BeamSearchNode node = head.getChild(0);
        while (node != null) {
            depth++;
            node = node.getChild(0);
        }
        return depth;
    }
    
    public List<BeamSearchBranch> getAllBranches() {
        return head.getAllBranches();
    }
    
    public List<BeamSearchBranch> getAllBranches(boolean includeHead) {
        return head.getAllBranches(includeHead);
    }
    
    public List<TokenList> getAllBranchTokensList() {
        List<BeamSearchBranch> allBranches = getAllBranches();
        List<TokenList> allBranchTokensList = new ArrayList<>();
        for (BeamSearchBranch branch : allBranches) {
            allBranchTokensList.add(branch.getTokenList());
        }
        return allBranchTokensList;
    }
}

